#!/usr/bin/env python
#-*-coding:utf-8-*-
import rospy
from sensor_msgs.msg import PointCloud2
import sensor_msgs.point_cloud2 as pc2
import numpy as np
import rospy
import cv2
import time
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
from message_filters import ApproximateTimeSynchronizer, Subscriber
import os
import os.path as osp
from nav_msgs.msg import Path,Odometry
import pdb
import math

def savePCD_bin_dim4(savename,point_cloud):
    point_cloud = point_cloud.astype(np.float32)
    num = point_cloud.shape[0]
    with open(savename, 'wb') as file:
        file.write(b'# .PCD v0.7 - Point Cloud Data file format\n')
        file.write(b'VERSION 0.7\n')
        file.write(b'FIELDS x y z intensity\n')
        file.write(b'SIZE 4 4 4 4\n')
        file.write(b'TYPE F F F F\n')
        file.write(b'COUNT 1 1 1 1\n')
        tt = 'WIDTH '+str(num)+'\n'
        file.write(tt.encode())
        file.write(b'HEIGHT 1\n')
        file.write(b'VIEWPOINT 0 0 0 1 0 0 0\n')
        dd = 'POINTS '+str(num)+'\n'
        file.write(dd.encode())
        file.write(b'DATA binary\n')
        point_cloud.tofile(file)

def mkdir_or_exist(dir_name, mode=0o777):
    if osp.exists(dir_name):
        return
    else:
        if dir_name == '':
            return
        dir_name = osp.expanduser(dir_name)
        os.makedirs(dir_name, mode=mode)

def getT(txt_path):
    with open(txt_path, 'r') as f:
        lines = f.readlines()

    data = []
    for line in lines:
        row = [float(x) for x in line.split()]
        data.append(row)

    data_array = np.array(data)
    data_array = np.concatenate((data_array,np.array([[0,0,0,1]])),axis=0)
    return data_array


def tf_points_new(xyz_points1,T1):
    xyz_points1_intensity = xyz_points1[:,3]
    xyz_points_hom = np.hstack((xyz_points1[:,:3], np.ones((xyz_points1.shape[0], 1))))
    xyz_points_transformed_hom = np.dot(T1, xyz_points_hom.T).T
    xyz_points_transformed = xyz_points_transformed_hom[:, :3]
    xyz_points_4dim_gps = np.hstack([xyz_points_transformed,xyz_points1_intensity.reshape(-1,1)])
    return xyz_points_4dim_gps


def image_callback(lidarleft, lidarright, image_tr, image_tl,image_fm,image_fl,image_fr,image_rr,image_rl,image_s):
    point_cloud_left = np.array([[x,y,z,intensity] for x,y,z,intensity in pc2.read_points(lidarleft,field_names=("x", "y", "z","intensity"),skip_nans=True)])
    point_cloud_right = np.array([[x,y,z,intensity] for x,y,z,intensity in pc2.read_points(lidarright,field_names=("x", "y", "z","intensity"),skip_nans=True)])

    bridge = CvBridge()
    im_tr = bridge.imgmsg_to_cv2(image_tr, desired_encoding='bgr8')
    im_tl = bridge.imgmsg_to_cv2(image_tl, desired_encoding='bgr8')
    im_fm = bridge.imgmsg_to_cv2(image_fm, desired_encoding='bgr8')
    im_fl = bridge.imgmsg_to_cv2(image_fl, desired_encoding='bgr8')
    im_fr = bridge.imgmsg_to_cv2(image_fr, desired_encoding='bgr8')
    im_rr = bridge.imgmsg_to_cv2(image_rr, desired_encoding='bgr8')
    im_rl = bridge.imgmsg_to_cv2(image_rl, desired_encoding='bgr8')
    im_s = bridge.imgmsg_to_cv2(image_s, desired_encoding='bgr8')

    time_stamp = str(lidarleft.header.stamp.secs)+str(lidarleft.header.stamp.nsecs)

    camera_type = ["frontleft","frontmid","frontright","rearleft","rearright","topleft","topright","stereo"]
    for cam,im in zip(camera_type,[im_fl,im_fm,im_fr,im_rl,im_rr,im_tl,im_tr,im_s]):
        print(osp.join(img_save_dir,cam,time_stamp+".jpg"))
        cv2.imwrite(osp.join(img_save_dir,cam,time_stamp+".jpg"),im)



    # save fusion img in one  jpg
    img_black = np.zeros_like(im_rl)
    set1 = np.hstack([img_black,im_tl,im_tr])
    set2 = np.hstack([im_fl,im_fm,im_fr])
    set3 = np.hstack([im_rl,img_black,im_rr])
    final = np.vstack([set1,set2,set3])
    cv2.imwrite(osp.join(img_save_dir,time_stamp+"final.jpg"), final)

    # save lidar
    lidar_l_E = getT(lidar_E_p+"/front_left.txt")
    lidar_r_E = getT(lidar_E_p+"/front_right.txt")
    points_l=tf_points_new(point_cloud_left,lidar_l_E)
    points_r=tf_points_new(point_cloud_right,lidar_r_E)

    pcd_l_save_p = os.path.join(save_dir, "lidar_left_ego",time_stamp+".pcd")
    mkdir_or_exist(osp.dirname(pcd_l_save_p))
    pcd_r_save_p = os.path.join(save_dir, "lidar_right_ego",time_stamp+".pcd")
    mkdir_or_exist(osp.dirname(pcd_r_save_p))
    pcd_fusion_save_p = os.path.join(save_dir, "lidar",time_stamp+".pcd")
    mkdir_or_exist(osp.dirname(pcd_fusion_save_p))
    #save org
    mkdir_or_exist(osp.dirname(osp.join(save_dir,"lidar_org",time_stamp+"_left.pcd")))
    savePCD_bin_dim4(osp.join(save_dir,"lidar_org",time_stamp+"_left.pcd"),point_cloud_left)
    savePCD_bin_dim4(osp.join(save_dir,"lidar_org",time_stamp+"_right.pcd"),point_cloud_right)
    
    # save ego
    savePCD_bin_dim4(pcd_l_save_p,points_l)
    savePCD_bin_dim4(pcd_r_save_p,points_r)
    savePCD_bin_dim4(pcd_fusion_save_p,np.vstack([points_l,points_r]))
    

    print("save ",osp.join(img_save_dir,time_stamp+"final.jpg"))

    exit()

import sys
if __name__ == "__main__":

    
    qtruck_num = sys.argv[1]
    cam_type= sys.argv[2]
    save_dir = sys.argv[3]+"/"+qtruck_num+"_"+cam_type

    img_save_dir = save_dir+"/image/"
    print(img_save_dir)
    mkdir_or_exist(img_save_dir)

    camera_type = ["frontleft","frontmid","frontright","rearleft","rearright","topleft","topright","stereo"]
    # camera_type = ["frontleft","frontmid","frontright","rearleft","rearright","topleft","topright"]
    for cam in camera_type:
        mkdir_or_exist(osp.join(img_save_dir,cam))

    lidar_E_p  = "/home/westwell/welldriver/wellpilot_config/calib/lidar_calib/"+qtruck_num

    # location_mat_dir = save_dir+"/location_mat"
    # mkdir_or_exist(location_mat_dir)

    rospy.init_node('sync_save')
    lidarLeft_sub = Subscriber('/lidar/left', PointCloud2)
    lidarRight_sub = Subscriber('/lidar/right', PointCloud2)
    image_sub_tr = Subscriber("/camera/top_right", Image)
    image_sub_tl = Subscriber("/camera/top_left", Image)

    image_sub_fm = Subscriber("/camera/front_mid", Image)
    image_sub_fl = Subscriber("/camera/front_left", Image)
    image_sub_fr = Subscriber("/camera/front_right", Image)
    image_sub_rr = Subscriber("/camera/rear_right", Image)
    image_sub_rl = Subscriber("/camera/rear_left", Image)
    image_sub_stereo = Subscriber("/stereo_image", Image)

    # image_sub_fm = Subscriber("/camera/top_right", Image)
    # image_sub_fl = Subscriber("/camera/top_right", Image)
    # image_sub_fr = Subscriber("/camera/top_right", Image)
    # image_sub_rr = Subscriber("/camera/top_right", Image)
    # image_sub_rl = Subscriber("/camera/top_right", Image)
    # image_sub_stereo = Subscriber("/camera/top_right", Image)

    
    # fusion_odom_sub = Subscriber('/fusion_odom', Odometry)
    ts = ApproximateTimeSynchronizer([lidarLeft_sub, lidarRight_sub,image_sub_tr, image_sub_tl,image_sub_fm,image_sub_fl,image_sub_fr,image_sub_rr,image_sub_rl,image_sub_stereo], queue_size=50, slop=0.5)
    ts.registerCallback(image_callback)

    rospy.spin()
